
import joblib  # 保存和加载模型
import numpy as np
import pandas as pd
import xgboost as xgb  # 极限梯度提升树对象
import matplotlib as plt
from sklearn.metrics import roc_auc_score
from xgboost import sklearn
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
import matplotlib.pyplot as plt

# 训练模型并获取特征重要性



df = pd.read_csv('../../data/raw/train.csv')


# 职业到管理级别的映射
job_to_level = {
    'Manufacturing Director': '高层管理',
    'Research Director': '高层管理',
    'Manager': '高层管理',
    'Sales Executive': '中层管理',
    'Laboratory Technician': '基层岗位',
    'Sales Representative': '基层岗位',
    'Healthcare Representative': '基层岗位',
    'Research Scientist': '基层岗位',
    'Human Resources': '基层岗位'
}
# 创建管理级别列
df['ManagementLevel'] = df['JobRole'].map(job_to_level)
level_to_code = {
    '基层岗位': 0,  # 基层编码为0
    '中层管理': 1,  # 中层编码为1
    '高层管理': 2  # 高层编码为2
}
# 直接映射为数值编码
df['JobRole'] = df['ManagementLevel'].map(level_to_code)
# 查看编码结果
# print(df[['JobRole', 'ManagementLevel', 'ManagementLevel_Encoded']].head())

df = pd.get_dummies(df,columns=['OverTime'])

x = df[['Age','DistanceFromHome','EmployeeNumber','MonthlyIncome','PercentSalaryHike','StockOptionLevel','TotalWorkingYears','OverTime_Yes','JobRole']]
y = df['Attrition']



def fit():
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=25)

    es = xgb.XGBRegressor(n_estimators=100, max_depth=2, learning_rate=0.1)

    es.fit(x_train, y_train)
    y_pre = es.predict(x_test)

    print(roc_auc_score(y_test, y_pre))

if __name__ == '__main__':
    fit()